import numpy as np
import torch
import csv
import torch.nn.functional as F
from metrics.registry import METRIC
from metrics import BaseMetric
from utils import get_logger
from torchmetrics.functional import accuracy

logger = get_logger("paddlevideo")


@METRIC.register
class SkeletonMetric(BaseMetric):
    """
    Test for Skeleton based model.
    note: only support batch size = 1, single card test.

    Args:
        out_file: str, file to save test results.
    """
    def __init__(self,
                 data_size,
                 batch_size,
                 out_file='submission.csv',
                 log_interval=50):
        """prepare for metrics
        """
        super().__init__(data_size, batch_size, log_interval)
        self.top1 = []
        self.top5 = []
        self.values = []
        self.out_file = out_file

    def update(self, batch_id, data, outputs):
        """update metrics during each iter
        """

        if len(data) == 2:  # data with label
            labels = data[1]
            top1 = accuracy(preds=outputs, target=labels, top_k=1)
            top5 = accuracy(preds=outputs, target=labels, top_k=5)
            self.top1.append(self.cast_numpy(top1))
            self.top5.append(self.cast_numpy(top5))
        else:  # data without label, only support batch_size=1. Used for fsd-10.
            prob = F.softmax(outputs,dim=-1)
            clas = torch.argmax(prob, dim=1)
            self.values.append((batch_id, self.cast_numpy(clas)))

        # preds ensemble
        if batch_id % self.log_interval == 0:
            logger.info("[TEST] Processing batch {}/{} ...".format(
                batch_id,
                self.data_size // self.batch_size))

    def cast_numpy(self,data):
        if isinstance(data,torch.Tensor):
            if data.is_cuda:
                data = data.cpu().detach().numpy()
            else:
                data = data.detach().numpy()
        return data

    def accumulate(self):
        """accumulate metrics when finished all iters.
        """
        if self.top1:  # data with label
            logger.info('[TEST] finished, avg_acc1= {}, avg_acc5= {}'.format(
                np.mean(np.array(self.top1)), np.mean(np.array(self.top5))))
        else:
            headers = ['sample_index', 'predict_category']
            with open(
                    self.out_file,
                    'w',
            ) as fp:
                writer = csv.writer(fp)
                writer.writerow(headers)
                writer.writerows(self.values)
            logger.info("Results saved in {} !".format(self.out_file))

